import torch
import pandas as pd

def pd_rolling_quantile(x: pd.DataFrame, window, quantile):
    return x.rolling(window, min_periods=1).quantile(quantile)

def torch_rolling_quantile(x: torch.Tensor, window: int, quantile: float, max_mem=4096) -> torch.Tensor:
    """
    计算二维张量 x 沿第 0 维的滑动分位数，窗口大小为 window。
    为满足显存限制，采用分块方式进行计算：
      1. 根据 max_mem (单位 MB) 计算每个批次允许的最大行数。
      2. 对每个批次构造滑动窗口数据，并计算窗口内的分位数。
      
    参数：
        x: 形状为 (N, M) 的二维张量
        window: 滑动窗口大小
        quantile: 分位数（介于 [0, 1] 之间）
        max_mem: 最大显存限制（单位 MB) 默认 4096
        
    返回：
        roll_quantile: 形状为 (N, M) 的张量，包含每个时刻的滑动分位数
    """
    with torch.no_grad():
        assert x.ndim == 2, "x 必须为二维张量"
        N, M = x.shape
        
        # 将非有限值（NaN 或 Inf）置为 NaN
        x = torch.where(torch.isfinite(x), x, torch.tensor(float('nan'), dtype=x.dtype, device=x.device))
        
        # 计算每个元素占用的字节数，并转换显存限制为字节
        element_size = x.element_size()
        max_mem_bytes = max_mem * 1024 * 1024  # MB 转字节
        
        # 计算当前批次内 rolled 数组的内存需求为：chunk_size * window * M * element_size
        # 从而确定每个批次允许处理的最大行数 chunk_size
        chunk_size = max_mem_bytes // (window * M * element_size) // 14
        chunk_size = max(1, chunk_size)  # 至少处理一行
        
        # 初始化结果张量（全部填充 NaN）
        roll_quantile = torch.full((N, M), float('nan'), dtype=x.dtype, device=x.device)
        
        # 分块处理每一部分数据
        for start in range(0, N, chunk_size):
            end = min(start + chunk_size, N)
            # 对当前批次构造时刻索引 (例如：[start, start+1, ..., end-1])
            current_indices = torch.arange(start, end, device=x.device)  # shape: (chunk_length,)
            offsets = torch.arange(-window + 1, 1, device=x.device)        # shape: (window,)
            # 对每个时刻，窗口索引为 [i-window+1, ..., i]
            indices = current_indices.unsqueeze(1) + offsets.unsqueeze(0)  # shape: (chunk_length, window)
            
            # 标记出有效索引（>=0），无效的部分后续填充 NaN
            valid_mask = indices >= 0
            indices_valid = indices.clamp(min=0)
            
            # 利用 advanced indexing 构造滑动窗口数据，结果形状为 (chunk_length, window, M)
            rolled = x[indices_valid]
            # 对无效的窗口位置填充 NaN
            rolled = rolled.masked_fill(~valid_mask.unsqueeze(-1), float('nan'))
            
            # 沿窗口维度 (dim=1) 计算分位数，torch.nanquantile 会忽略 NaN 值
            quantile_vals = torch.nanquantile(rolled, quantile, dim=1)
            quantile_vals[~torch.isfinite(quantile_vals)] = float('nan')
            
            # 将当前批次结果填入最终结果张量中
            roll_quantile[start:end] = quantile_vals
        
        return roll_quantile
